package interceptor

import (
	"context"
	"go-rpc-2/server_register_and_found/register"
	"google.golang.org/grpc"
	"reflect"
	"sync"
)

type broadcast struct {
}

func SetBroadcast(ctx context.Context, ch chan Resp) context.Context {
	return context.WithValue(ctx, broadcast{}, ch)
}

func isBroadcast(ctx context.Context) (chan Resp, bool) {
	ch := ctx.Value(broadcast{})
	if ch != nil {
		return ch.(chan Resp), true
	}
	return nil, false
}

type BroadcastBuilder struct {
	serviceName string
	register    register.Register
	opts        []grpc.DialOption
}

func NewBroadcastBuilder(serviceName string, register2 register.Register, opts ...grpc.DialOption) *BroadcastBuilder {
	return &BroadcastBuilder{serviceName: serviceName, register: register2, opts: opts}
}

func (b *BroadcastBuilder) Build() grpc.UnaryClientInterceptor {
	return func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
		ch, ok := isBroadcast(ctx)
		if !ok {
			return invoker(ctx, method, req, reply, cc, opts...)
		}

		si, err := b.register.ListService(ctx, b.serviceName)
		if err != nil {
			return err
		}
		wg := sync.WaitGroup{}
		wg.Add(len(si))
		typ := reflect.TypeOf(reply).Elem()

		for _, service := range si {
			go func() {
				defer wg.Done()
				var siConn *grpc.ClientConn
				resp := reflect.New(typ).Interface()
				siConn, err = grpc.Dial(service.Addr, b.opts...)
				if err != nil {
					ch <- Resp{Err: err}
					return
				}

				err = invoker(ctx, method, req, resp, siConn, opts...)
				if err != nil {
					ch <- Resp{Err: err}
					return
				}

				//select { // 返回第一个响应
				//case ch <- Resp{
				//	Data: resp,
				//}:
				//default:
				//
				//}

				select {
				case <-ctx.Done():
					ch <- Resp{Err: ctx.Err()}
					return
				case ch <- Resp{Data: resp}:
				}
			}()
		}

		wg.Wait()
		return nil

	}
}

type Resp struct {
	Data any
	Err  error
}
